from typing import Union
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.hgcn_api_service import HGCNService
from app.schemas.dhgam.hgcn_schema import HGCNRequest, HGCNResponse, HGCNMiddleResponse
from app.schemas.facenet_base_schema import TaskStatus
from app.tools.redis_client import redis_client
import logging

# 配置日志
logger = logging.getLogger(__name__)
router = APIRouter()
service = HGCNService()
# 异构图卷积神经网络 (HGCN) 算法简介：能够处理包含不同类型节点和边的图结构数据，为每个节点生成低维向量表示（嵌入）。
@router.post("/create", response_model=HGCNMiddleResponse)
async def create_task(request: HGCNRequest) -> HGCNMiddleResponse:
    """
    创建新的HGCN任务并存储在Redis中
    
    Args:
        request: 包含任务ID、回调URL和输入参数的请求对象
        
    Returns:
        HGCNMiddleResponse: 包含任务初始状态的响应对象
    """
    try:
        task_response = HGCNMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.dict())
        return task_response
    except Exception as e:
        logger.error(f"创建HGCN任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"创建任务失败: {str(e)}")

@router.post("/execute/{task_id}", response_model=HGCNMiddleResponse)
async def execute_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> HGCNMiddleResponse:
    """
    执行指定ID的HGCN任务
    
    Args:
        task_id: 任务唯一标识
        background_tasks: FastAPI的后台任务对象
        
    Returns:
        HGCNMiddleResponse: 任务执行状态的响应对象
    """
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="任务不存在")

        input_params = task_data.get("input_params")
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        background_tasks.add_task(run_task_in_background, input_params, task_id)
        
        return HGCNMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行HGCN任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"执行任务失败: {str(e)}")

@router.get("/result/{task_id}", response_model=Union[HGCNResponse, HGCNMiddleResponse])
async def get_result(task_id: str) -> Union[HGCNResponse, HGCNMiddleResponse]:
    """
    获取HGCN任务执行结果
    
    Args:
        task_id: 任务唯一标识
        
    Returns:
        Union[HGCNResponse, HGCNMiddleResponse]: 
            - 如果任务完成，返回HGCNResponse
            - 如果任务进行中，返回HGCNMiddleResponse
    """
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="结果不存在")
        
        # 确保result_data包含所有必需的字段
        if "input_params" not in result_data:
            # 尝试从原始任务数据中获取input_params
            original_task_data = await redis_client.get_data(task_id)
            if original_task_data and "input_params" in original_task_data:
                result_data["input_params"] = original_task_data["input_params"]
            else:
                # 如果无法获取，则设置为空字典避免验证错误
                result_data["input_params"] = {}
        
        return (HGCNResponse(**result_data) 
                if "output_params" in result_data 
                else HGCNMiddleResponse(**result_data))
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取HGCN结果失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"获取结果失败: {str(e)}")

async def run_task_in_background(input_params, task_id: str) -> None:
    """
    在后台执行HGCN算法任务
    
    Args:
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    try:
        result = await service.process_hgcn(input_params, task_id)
        # 确保结果中包含input_params
        if "input_params" not in result:
            result["input_params"] = input_params
        await redis_client.update_data(f'{task_id}_result', result)
    except Exception as e:
        logger.error(f"HGCN后台任务执行失败: {str(e)}", exc_info=True)
        # 更新任务状态为失败
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e),
            "input_params": input_params  # 确保错误状态也包含input_params
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 